import os
import matplotlib
matplotlib.use('Agg')  # 设置为非交互式后端，避免Tkinter线程问题
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd


# 可视化工具类
class TrainingVisualizer:
    def __init__(self, log_dir='visualizations', save_interval=1):
        self.log_dir = log_dir
        self.save_interval = save_interval
        os.makedirs(self.log_dir, exist_ok=True)
        
        self.epochs = []
        self.train_losses = []
        self.contrastive_losses = []
        self.feature_matching_losses = []
        self.alpha_values = []
        self.eval_metrics = {}
        
    def update_train_loss(self, epoch, loss, contrastive_loss=None, feature_matching_loss=None, alpha=None):
        self.epochs.append(epoch)
        self.train_losses.append(loss)
        if contrastive_loss is not None:
            self.contrastive_losses.append(contrastive_loss)
        if feature_matching_loss is not None:
            self.feature_matching_losses.append(feature_matching_loss)
        if alpha is not None:
            self.alpha_values.append(alpha)
        
    def update_eval_metrics(self, epoch, metrics):
        for metric_name, value in metrics.items():
            if metric_name not in self.eval_metrics:
                self.eval_metrics[metric_name] = []
            self.eval_metrics[metric_name].append(value)
        
    def plot_all(self, epoch):
        self._plot_train_loss()
        self._plot_eval_metrics()
        
        if epoch % self.save_interval == 0:
            self._save_to_csv()
    
    def _plot_train_loss(self):
        plt.figure(figsize=(12, 8))
        plt.plot(self.epochs, self.train_losses, 'b-', label='Total Loss')
        if self.contrastive_losses:
            plt.plot(self.epochs, self.contrastive_losses, 'g-', label='Contrastive Loss')
        if self.feature_matching_losses:
            plt.plot(self.epochs, self.feature_matching_losses, 'r-', label='Feature Matching Loss')
        plt.title('Training Losses')
        plt.xlabel('Epoch')
        plt.ylabel('Loss')
        plt.grid(True)
        plt.legend()
        plt.tight_layout()
        plt.savefig(os.path.join(self.log_dir, 'train_loss.png'))
        plt.close()
        
        # 绘制alpha值变化
        if self.alpha_values:
            plt.figure(figsize=(10, 4))
            plt.plot(self.epochs, self.alpha_values, 'm-', label='Alpha Value')
            plt.title('Alpha Value Changes')
            plt.xlabel('Epoch')
            plt.ylabel('Alpha')
            plt.ylim(0, 1)
            plt.grid(True)
            plt.legend()
            plt.tight_layout()
            plt.savefig(os.path.join(self.log_dir, 'alpha_changes.png'))
            plt.close()
    

    
    def _plot_eval_metrics(self):
        if not self.eval_metrics:
            return
        
        plt.figure(figsize=(12, 8))
        for metric_name, values in self.eval_metrics.items():
            if len(values) == len(self.epochs):
                plt.plot(self.epochs, values, marker='o', label=metric_name)
        
        plt.title('评估指标曲线')
        plt.xlabel('Epoch')
        plt.ylabel('指标值')
        plt.grid(True)
        plt.legend()
        plt.tight_layout()
        plt.savefig(os.path.join(self.log_dir, 'eval_metrics.png'))
        plt.close()
    
    def _save_to_csv(self):
        data = {'epoch': self.epochs, 'train_loss': self.train_losses}
        if self.contrastive_losses:
            data['contrastive_loss'] = self.contrastive_losses
        if self.feature_matching_losses:
            data['feature_matching_loss'] = self.feature_matching_losses
        if self.alpha_values:
            data['alpha'] = self.alpha_values
        for metric_name, values in self.eval_metrics.items():
            if len(values) == len(self.epochs):
                data[metric_name] = values
            else:
                padded_values = [np.nan] * len(self.epochs)
                padded_values[:len(values)] = values
                data[metric_name] = padded_values
        
        df = pd.DataFrame(data)
        df.to_csv(os.path.join(self.log_dir, 'training_history.csv'), index=False)
